"""
@File    : rsio.py.py
@Time    : 2025/7/15 下午5:21
@Author  : xiashuobad
@Desc    : 

/**
 * _ooOoo_
 * o8888888o
 * 88" . "88
 * (| -_- |)
 *  O\ = /O
 * ___/`---'\____
 * .   ' \\| |// `.
 * / \\||| : |||// \
 * / _||||| -:- |||||- \
 * | | \\\ - /// | |
 * | \_| ''\---/'' | |
 * \ .-\__ `-` ___/-. /
 * ___`. .' /--.--\ `. . __
 * ."" '< `.___\_<|>_/___.' >'"".
 * | | : `- \`.;`\ _ /`;.`/ - ` : | |
 * \ \ `-. \_ __\ /__ _/ .-` / /
 * ======`-.____`-.___\_____/___.-`____.-'======
 * `=---=' bug泛滥 佛已瘫痪
"""
# rsio/rsio.py
import os
import logging
import numpy as np
import math

from functools import reduce
from pathlib import Path
from typing import Protocol, Union, Type, overload, Literal, Sequence, IO, Any, Generator
from itertools import product
from osgeo import osr, gdal
from importlib import util

from .backends import ImageBackend
from .utils import RichLogger, ArrayShapeConverter

logger = RichLogger(logging.getLogger(__name__))
gdal.SetCacheMax(0)
if util.find_spec('rasterio'):
    import rasterio

    rasterio.Env(GDAL_CACHEMAX=0).__enter__()


class ImageBackendProtocol(Protocol):
    """定义图像后端的公共接口，用于类型提示"""

    # 基本属性

    @property
    def img_data(self):
        ...

    def set_img_data(self, img_data) -> 'ImageBackendProtocol':
        ...

    @property
    def img_array(self) -> np.ndarray:
        """
        获取图像的array数据
        :return:
        """
        ...

    def set_img_array(self, array: np.ndarray) -> 'ImageBackendProtocol':
        """
        设置图像array
        :param array: numpy.ndarray
        :return: 图像对象
        """
        ...

    @property
    def palette(self) -> list:
        ...

    def set_palette(self, palette: Union[list, tuple, dict, bytes, np.ndarray]) -> 'ImageBackendProtocol':
        ...

    @property
    def shape(self) -> tuple[int, int] | tuple[int, int, int]:
        """
        返回图像的array的shape
        :return:
        """
        ...

    @property
    def size(self) -> tuple[int, int]:
        """
        返回图像的array的size：(宽,高)
        :return:
        """
        ...

    @property
    def channel(self) -> int:
        """
        返回图像的array的通道数
        :return:
        """
        ...

    @property
    def width(self) -> int:
        """
        返回图像的array的宽度
        :return:
        """
        ...

    @property
    def height(self) -> int:
        """
        返回图像的array的高度
        :return:
        """
        ...

    @property
    def dtype(self) -> np.dtype | str | Type[np.number]:
        """
        返回图像array的数据类型
        :return:
        """
        ...

    @property
    def file_path(self) -> str:
        """
        返回图像文件的路径
        :return:
        """
        ...

    @property
    def mode(self) -> str:
        """
        返回图像的打开模式
        :return:
        """
        ...

    @property
    def kwargs(self) -> dict:
        """
        返回图像的附加打开参数
        :return:
        """
        ...

    @property
    def file_name(self) -> str:
        """
        返回图像的文件名称
        @return:
        """
        ...

    @property
    def img_suffix(self) -> str:
        """
        返回图像的文件后缀
        @return:
        """
        ...

    @property
    def geo_bounds(self) -> tuple[float, float, float, float]:
        """
        返回图像的地理边界
        :return:
        """
        ...

    @property
    def nodata(self) -> int | float:
        ...

    def set_nodata(self, nodata: int | float) -> 'ImageBackendProtocol':
        ...

    @property
    def geo_transform(self) -> 'tuple[float, ...] | affine.Affine':
        ...

    def set_geo_transform(self, geo_transform: 'tuple[float, ...] | affine.affine') -> 'ImageBackendProtocol':
        ...

    @property
    def geo_projection(self) -> str:
        ...

    def set_geo_projection(self, geo_projection: str | int | osr.SpatialReference) -> 'ImageBackendProtocol':
        ...

    def read_window_array(
            self,
            x_off: int = 0,
            y_off: int = 0,
            window_x_size: int = None,
            window_y_size: int = None,
            band_indexes: int | list[int] = None,
            **kwargs
    ) -> np.ndarray:
        """
        读取窗口区域数据
        :param x_off: 窗口左上角X坐标
        :param y_off: 窗口左上角Y坐标
        :param window_x_size: 窗口宽度
        :param window_y_size: 窗口高度
        :param band_indexes: 读取的波段索引，默认读取所有波段，使用list可以任意组合波段，例如[2,1,0]可以读取前三个波段并反转rgb
        :param kwargs: 附加参数，根据不同后端设置
        :return: numpy.ndarray
        """
        ...

    def save(self, save_path=None, **meta):
        """
        将图像数据保存到文件
        :param save_path: 重设保存路径
        :param meta:
        :return:
        """
        ...

    def resize(self,
               target_size: tuple[int, int] | int,
               keep_ratio: bool = False,
               pad: int = None,
               resample: int = None,
               box: list[int, int, int, int] | tuple[int, int, int, int] = None,
               **kwargs
               ) -> 'ImageBackendProtocol':
        """
        对图像进行resize操作
        :param target_size: 目标尺寸 (width,height)
        :param keep_ratio: 是否保持宽高比
        :param pad: 在保持宽高比时进行padding填充，pad表示要填充的值
        :param resample: 重采样的策略，需要根据不同的后端来设置
        :param box: 目标框：[x_min,y_min,x_max,y_max]，设置这个时会先进行目标框的裁剪，再对目标框的图像进行后续操作
        :param kwargs: 可选附加参数，根据不同后端设置
        :return: resize后的图像对象
        """
        ...

    def reverse_rgb(self) -> 'ImageBackendProtocol':
        """
        将图片的RGB通道进行反转，比如将RGB变为BGR
        :return: image实例对象
        """
        ...

    def crop(self,
             x_off: int = 0,
             y_off: int = 0,
             window_x_size: int = None,
             window_y_size: int = None
             ) -> 'ImageBackendProtocol':
        """
        裁剪图片
        :param x_off: 裁剪的起始x坐标
        :param y_off: 裁剪的起始y坐标
        :param window_x_size: 裁剪的宽度
        :param window_y_size: 裁剪的高度
        :return: 裁剪后的图片对象
        """
        ...

    def close(self):
        """
        关闭图片对象
        :return:
        """
        ...

    def show(self, window_name: str = None):
        """
        显示图片
        :param window_name: 窗口名称
        :return:
        """
        ...

    def reproject(self,
                  dst_srs_obj: 'int|str|ImageBackendProtocol',
                  dst_width=None,
                  dst_height=None,
                  num_threads: int = 2):
        """
        对图片进行重投影
        :param dst_srs_obj: 目标投影对象
        :param dst_width: 目标宽度
        :param dst_height: 目标高度
        :param num_threads: 重投影的线程数
        :return: 重投影后的图片对象
        """
        ...

    def convert_to_rgb(self):
        """
        将图片转换为RGB格式
        :return:
        """
        ...


class OpenImage:
    registered_backends = ImageBackend.registered_backends

    def __init__(self, fp: Union[str, Path], mode: str = 'r', backend: str = 'gdal', **kwargs):
        self.backend_image: ImageBackendProtocol = self.registered_backends[backend](fp, mode, **kwargs)
        self.backend_name = backend

    def __enter__(self) -> 'ImageBackendProtocol|OpenImage':
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.backend_image.close()

    def __getattr__(self, attr):
        return getattr(self.backend_image, attr)

    def set_backend_image(self, backend_image):
        self.backend_image = backend_image
        return self

    @property
    def pixel_size(self) -> tuple[float, float]:
        """
        像素大小/像元大小
        :return: (x_pixel_size,y_pixel_size)
        """
        try:
            x_pixel_size, y_pixel_size = self.geo_transform[1], abs(self.geo_transform[5])
        except Exception:
            x_pixel_size, y_pixel_size = 1, 1
        return x_pixel_size, y_pixel_size

    def sliding_window_x_y(
            self,
            window_size: int | tuple[int, int] = (512, 512),
            stride: int | tuple[int, int] | list[int] = None,
    ):
        if isinstance(window_size, int):
            window_size = (window_size, window_size)
        if len(window_size) == 1:
            window_size = window_size * 2
        if not stride:
            stride = window_size
        if isinstance(stride, int):
            stride = (stride, stride)
        list_x = list(range(0, self.width - window_size[0], stride[0])) + [self.width - window_size[0]]
        list_y = list(range(0, self.height - window_size[-1], stride[-1])) + [self.height - window_size[-1]]
        return product(list_x, list_y)

    def sliding_window_tiles(
            self,
            window_size: int | tuple[int, int] = (512, 512),
            stride: int | tuple[int, int] | list[int] = None,
            max_tiles: int = None,
            band_indexes: int | list[int] = None
    ):
        """
        滑动窗口对大图像进行切片，返回一个生成器
        @param window_size: 窗口大小, (width,height)
        @param stride: 步长, (x_stride,y_stride)
        @param max_tiles: 最大切片数量，默认切割所有，设置后会限制切片数量
        @param band_indexes: 波段索引，默认切割所有波段
        @return:
        """
        if isinstance(window_size, int):
            window_size = (window_size, window_size)
        if len(window_size) == 1:
            window_size = window_size * 2

        for count, (x, y) in enumerate(self.sliding_window_x_y(window_size, stride), start=1):
            yield self.backend_image.read_window_array(x, y, *window_size, band_indexes), (x, y)
            if max_tiles and count == max_tiles:
                break

    def count_tiles(
            self,
            stride: int | tuple[int, int] | list[int]
    ) -> int:
        """
        计算滑动窗口切片数量
        @param stride: 步长, (x_stride,y_stride)
        @return:
        """
        if isinstance(stride, int):
            stride = (stride, stride)
        stride_x, stride_y = stride[0], stride[-1]
        return math.ceil(self.width / stride_x) * math.ceil(self.height / stride_y)

    def reproject(self,
                  dst_srs_obj: 'int|str|OpenImage',
                  dst_width=None,
                  dst_height=None,
                  num_threads: int = 2) -> 'ImageBackendProtocol|OpenImage':
        """
        对图片进行重投影
        :param dst_srs_obj: 参考投影对象，支持int类型epsg代码 xxxx，字符串格式 "epsg:xxxx","EPSG:xxxx"，wkt字符串，或者OpenImage对象
        :param dst_width: 目标宽度
        :param dst_height: 目标高度
        :param num_threads: 重投影的线程数
        :return: 重投影后的图片对象
        """
        if isinstance(dst_srs_obj, OpenImage):
            dst_srs_obj = dst_srs_obj.backend_image
        self.backend_image.reproject(dst_srs_obj, dst_width, dst_height, num_threads)
        return self

    def valid_mask(self):
        """
        计算本图像的有效值掩码
        """
        return OpenImage.valid_mask_of_image(self.img_array)

    def image_blocks(self, num_blocks) -> Generator[tuple[ImageBackendProtocol, tuple[int, int]], None, None]:
        """
        图像均匀分块
        @param num_blocks: 分块数量
        @return: 返回所有的图像块
        """
        rows, cols = OpenImage.calculate_rows_cols_for_blocks(self.width, self.height, num_blocks)
        block_width = math.ceil(self.width / cols)
        block_height = math.ceil(self.height / rows)
        list_x = list(range(0, self.width - block_width, block_width)) + [self.width - block_width]
        list_y = list(range(0, self.height - block_height, block_height)) + [self.height - block_height]
        for x, y in product(list_x, list_y):
            image = self.crop(x, y, block_width, block_height)
            yield image, (x, y)

    def crop_by_geojson(self, geojson: dict | list[dict]):
        """
        根据地理区域裁剪图片
        :param geojson: 地理范围
        :return: 裁剪后的图片对象
        """
        assert isinstance(geojson, (dict, list)), "geojson参数类型错误！"
        if isinstance(geojson, dict):
            if 'features' in geojson:
                geojson = list(map(lambda feature: feature['geometry'], geojson['features']))
            elif 'geometry' in geojson:
                geojson = [geojson['geometry']]
            else:
                geojson = [geojson]
        assert self.backend_name == 'rasterio', "该功能需要使用rasterio后端！"
        from rasterio.mask import mask
        out_image, out_transform = mask(self.img_data, geojson, nodata=self.nodata, crop=True)
        new_image = open_image(
            self.file_path.replace(self.img_suffix, f'_by_geojson{self.img_suffix}'),
            'w',
            backend='rasterio'
        )
        new_image.set_img_array(out_image).set_geo_transform(out_transform).set_geo_projection(
            self.geo_projection).set_nodata(self.nodata)
        return new_image

    def crop_by_geo_bounds(
            self,
            geo_bounds: tuple[float, float, float, float],
            output_path: str = None
    ) -> 'OpenImage|ImageBackendProtocol':
        """
        根据地理范围裁剪图片
        :param geo_bounds: 地理范围 (左，上，右，下)
        :return: 裁剪后的图片对象
        """
        assert self.backend_name == 'gdal', "该功能需要使用gdal后端！"
        if output_path is None:
            output_path, _format = '', 'MEM'
        else:
            _format = 'GTiff'
        cropped_ds: gdal.Dataset = gdal.Translate(
            output_path,
            self.img_data,
            projWin=list(geo_bounds),
            format=_format,
        )
        if not cropped_ds:
            logger.exception(error := '根据地理范围裁剪图片失败！')
            raise Exception(error)
        if output_path:
            cropped_ds = None
            new_image = self.__class__(output_path)
        else:
            new_image = self.__class__(self.file_path, 'w')
            new_image.set_img_data(cropped_ds)
        return new_image

    @overload
    @staticmethod
    def write_array(
            file_path: str,
            array: np.ndarray,
            backend: Literal['gdal'] = 'gdal',
            *,
            geo_transform: tuple[float] = None,
            geo_projection: str = None,
            nodata: int | float = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray],
            **kwargs
    ):
        """
        将指定数组array写入文件，后端为gdal时使用
        :param file_path: 文件路径
        :param array: 数组
        :param backend: 后端为gdal
        :param geo_transform: 地理坐标变换参数
        :param geo_projection: 地理投影
        :param nodata: 缺失值
        :param palette: 颜色表
        :param kwargs: 其他参数，用于 gdal.dataset.WriteArray(array,**kwargs)
        :return:
        """
        ...

    @overload
    @staticmethod
    def write_array(file_path: str,
                    array: np.ndarray,
                    backend: Literal['cv2'] = 'cv2',
                    *,
                    params: Sequence[int] = None
                    ):
        """
        将指定数组array写入文件，后端为cv2时使用
        :param file_path: 文件路径
        :param array: 数组
        :param backend: 后端为cv2
        :param params: 其他参数，用于 cv2.imencode(ext, array, params)
        :return:
        """
        ...

    @overload
    @staticmethod
    def write_array(
            fp: Union[str, bytes, os.PathLike[str], os.PathLike[bytes]] | IO[bytes],
            array: np.ndarray,
            backend: Literal['pil'] = 'pil',
            *,
            mode: str = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            file_format: str = None,
            **params: Any
    ):
        """
        将指定数组array写入文件，后端为pil时使用
        :param fp:
        :param array:
        :param backend:
        :param mode: 图片模式，与pil中的mode相同
        :param palette: 颜色映射，用于PIL.Image.putpalette(palette)
        :param file_format: 用于 PIL.Image.save(fp, file_format, **params)
        :param params: 用于 PIL.Image.save(fp, file_format, **params)
        :return:
        """
        ...

    @overload
    @staticmethod
    def write_array(
            fp: str,
            array: np.ndarray,
            backend: Literal['rasterio'],
            *,
            geo_transform: tuple[float, ...] = None,
            geo_projection: str = None,
            nodata: int | float = None,
            palette: Union[list, tuple, dict, bytes, np.ndarray] = None,
            **meta
    ):
        """
        将指定数组array写入文件，后端为rasterio时使用
        :param fp:
        :param array:
        :param backend:
        :param geo_transform:
        :param geo_projection:
        :param nodata:
        :param palette:
        :param meta: 图像元信息，用于 rasterio.open(fp, 'w', **meta)
        :return:
        """
        ...

    @staticmethod
    def write_array(file_path: str | Path, array: np.ndarray, backend: str = 'gdal', **kwargs):
        OpenImage.registered_backends[backend].write_array(file_path, array, **kwargs)

    @staticmethod
    def show_image(img_array: np.ndarray, window_name: str = 'image', backend: str = 'pil'):
        """
        显示指定图像array
        :param img_array:
        :param window_name:
        :param backend:
        :return:
        """
        OpenImage.registered_backends[backend].show_image(img_array, window_name)

    @staticmethod
    def resize_img(
            img_array: np.ndarray,
            target_size: tuple[int, int] | int,
            keep_ratio: bool = False,
            pad: int = None,
            out: np.ndarray = None,
            fx: float = None,
            fy: float = None,
            resample: int = None,
            box: list[int] | tuple[int, int, int, int] = None,
            shape_style: str = 'cv2',
    ) -> np.ndarray:
        """
        对指定图像array进行resize
        :param img_array: 图像array
        :param target_size: (width,height)
        :param keep_ratio: 是否保持宽高比
        :param pad: keep_ratio为True时进行pad填充
        :param out:结果赋值到已存在的变量内存上
        :param fx: width缩放系数
        :param fy: height缩放系数
        :param resample: 重采样策略，同cv2
        :param box: 指定窗口resize，[x_min,y_min,x_max,y_max]
        :param shape_style:同后端名称，输出结果的shape格式,不同后端格式有区别: (c,h,w) 或者 (h,w,c)
        :return: resize后的array
        """
        return OpenImage.registered_backends['cv2'].resize_img(
            img_array,
            target_size,
            keep_ratio,
            pad,
            out,
            fx,
            fy,
            resample,
            box,
            shape_style
        )

    @staticmethod
    def reversed_rgb(img_array: np.ndarray, shape_style: str = 'cv2') -> np.ndarray:
        """
        :param img_array: numpy数组
        :param shape_style: array shape的风格，['gdal','cv2']
        :return: 反转后的数组
        """
        return OpenImage.registered_backends['cv2'].reversed_rgb(img_array, shape_style)

    @staticmethod
    def intersection_bounds(bounds_list: list[tuple[float, float, float, float]]):
        """
        计算多个bounds的公共区域的边界框bounds
        bounds: (left,top,right,bottom)
        @param bounds_list: [bounds1,bounds2,...]，
        @return:
        """

        # 初始化为第一个bounds
        intersection_left, intersection_top, intersection_right, intersection_bottom = bounds_list[0]

        # 依次与后续的bounds求交集
        for bounds in bounds_list[1:]:
            left, top, right, bottom = bounds

            # 更新交集边界
            intersection_left = max(intersection_left, left)
            intersection_top = min(intersection_top, top)  # top是较大的y值
            intersection_right = min(intersection_right, right)
            intersection_bottom = max(intersection_bottom, bottom)  # bottom是较小的y值

            # 检查是否还有交集
            if intersection_left >= intersection_right or intersection_top <= intersection_bottom:
                return None

        return intersection_left, intersection_top, intersection_right, intersection_bottom

    @staticmethod
    def intersection_image(image_list: list['OpenImage|ImageBackendProtocol']) -> list['OpenImage']:
        """
        求多个OpenImage对象的公共地理区域后并返回各自修改后的OpenImage
        @param image_list: OpenImage列表
        @return:
        """
        assert len(image_list) >= 2
        geo_projection = image_list[0].geo_projection
        for image in image_list[1:]:
            if not OpenImage.is_same_projection(geo_projection, image.geo_projection):
                logger.warning(f'{image.file_name} 与原图投影不一致，将进行投影转换...', markup='yellow')
                image.reproject(geo_projection)
        bounds_list = [image.geo_bounds for image in image_list]
        if any(not OpenImage.is_equal_bounds(bounds_list[0], bounds) for bounds in bounds_list[1:]):
            intersection_geo_bounds = OpenImage.intersection_bounds(bounds_list)
            if intersection_geo_bounds is None:
                raise ValueError("这些图像没有公共地理区域！")
            logger.info(f'截取图像的公共地理区域...', markup='green')
            for i, image in enumerate(image_list):
                image_list[i] = image.crop_by_geo_bounds(intersection_geo_bounds)
                image.close()
        return image_list

    @staticmethod
    def is_same_projection(proj1_str, proj2_str) -> bool:
        """
        使用GDAL比较两个投影字符串是否表示相同的坐标系
        """
        # 创建空间参考对象
        srs1 = osr.SpatialReference()
        srs2 = osr.SpatialReference()
        # 导入投影信息
        if srs1.ImportFromWkt(proj1_str) != 0:
            # 如果不是WKT格式，尝试其他格式
            srs1.ImportFromProj4(proj1_str)
        if srs2.ImportFromWkt(proj2_str) != 0:
            srs2.ImportFromProj4(proj2_str)
        # 标准化坐标系
        srs1.AutoIdentifyEPSG()
        srs2.AutoIdentifyEPSG()
        # 比较是否相同
        return srs1.IsSame(srs2)

    @staticmethod
    def is_equal_bounds(bounds1, bounds2, rtol=1e-9, atol=1e-12):
        """
        比较两个地理边界是否相等（考虑浮点数精度）

        Args:
            bounds1: 第一个边界 (min_x, max_y, max_x, min_y)
            bounds2: 第二个边界 (min_x, max_y, max_x, min_y)
            rtol: 相对容差
            atol: 绝对容差

        Returns:
            bool: 是否相等
        """
        return np.allclose(bounds1, bounds2, rtol=rtol, atol=atol)

    @staticmethod
    def valid_mask_of_image(img_array: np.ndarray):
        """
            计算所给图像的有效值掩码,即多个通道不全为0或者不全为255的为有效
            :param img_array:
            :return:
        """
        img_array = ArrayShapeConverter.array_shape_to_gdal(img_array)
        mask_nodata1 = img_array != 0
        mask_nodata2 = img_array != 255
        if img_array.ndim == 3:  # 如果是rgb影像，需要对多个通道进行判断有效值，并转为单通道掩码，与最终预测结果shape保持一致
            mask_nodata1 = reduce(lambda _x, _y: _x | _y, mask_nodata1)  # 多个通道取并，即只要任意一个通道不是无效值，就判断为有效
            mask_nodata2 = reduce(lambda _x, _y: _x | _y, mask_nodata2)
        return mask_nodata1 * mask_nodata2

    @staticmethod
    def calculate_rows_cols_for_blocks(image_width, image_height, num_blocks):
        """
        计算将图像分割成指定数量的面积相近的块的最优方案

        Args:
            image_width (int): 原始图像的宽度
            image_height (int): 原始图像的高度
            num_blocks (int): 要分割的块数

        Returns:
            tuple: (rows, cols) 表示在行和列上的分割数
        """
        if num_blocks <= 0:
            raise ValueError("num_blocks must be a positive integer")

        # 寻找最优的行列分割数
        best_rows, best_cols = 1, num_blocks
        best_diff = float('inf')

        # 遍历所有可能的因数组合
        for rows in range(1, int(math.sqrt(num_blocks)) + 1):
            if num_blocks % rows == 0:
                cols = num_blocks // rows
                # 计算每个块的尺寸
                block_width = image_width / cols
                block_height = image_height / rows
                # 计算宽高差值
                diff = abs(block_width - block_height)

                # 检查当前组合是否更优
                if diff < best_diff:
                    best_diff = diff
                    best_rows, best_cols = rows, cols

                # 也检查行列互换的情况
                block_width_alt = image_width / rows
                block_height_alt = image_height / cols
                diff_alt = abs(block_width_alt - block_height_alt)

                if diff_alt < best_diff:
                    best_diff = diff_alt
                    best_rows, best_cols = cols, rows

        return best_rows, best_cols


def open_image(file_path: str, mode: str = 'r', backend: str = 'gdal', **kwargs) -> ImageBackendProtocol | OpenImage:
    return OpenImage(file_path, mode, backend, **kwargs)
